#!py -3
# coding:utf-8

import os
import onnx
import torch
import torch.nn as nn
import numpy as np
import cv2 as cv
import netron
import torch.onnx
import onnxruntime
import numpy as np
from torchvision.transforms import transforms
from Unet import *
from data_pk import *

use_gpu = torch.cuda.is_available()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 定义训练使用的设备
if __name__ == '__main__':
    ###########新建数据集###############
    inputlist = getFilepathInDir("../oct_img/22")  # ("../oct_img/input")
    saveinavi = "../oct_img/input.avi"
    saveoutavie = "../oct_img/output.avi"

    reqs = (512, 512)

    dataset = ImgListDataSet(inputlist, reqsize=reqs)
    totensor = transforms.ToTensor()

    #########新建模型##############
    # model = UNet(in_channel=1, out_channel=1)
    # model = LessUNet(in_channel=1, out_channel=2)
    # model.add_module("non_linear_layer", MySigmoid())
    model = MixUNet1()
    model: torch.nn.Module
    model.to(device)

    #############加载模型参数#############
    # savemoduleparamName = "../moduleparam/enhance_unet_sigmoid.pth"
    # savemoduleparamName = "../moduleparam/enhance_lessunet2.pth"
    # savemoduleparamName = "../moduleparam/enhance_mixlessunet1.pth"
    savemoduleparamName = "../moduleparam/enhance_mixlessunet4_1.pth"
    if os.path.exists(savemoduleparamName):
        model.load_state_dict(torch.load(savemoduleparamName))
    else:
        newdir, _ = os.path.split(savemoduleparamName)
        os.makedirs(newdir, exist_ok=True)

    ##############导出为onnx模型######################
    # onnx_filename = '../moduleparam/MixUNet1.onnx'
    onnx_filename = '../moduleparam/MixUNet4_2.onnx'
    dynamic_axes = {
        'input': {0: 'batch', 2: 'height', 3: 'width'},
        'output': {0: 'batch', 2: 'height', 3: 'width'},
    }
    x = totensor(dataset[0])
    x: torch.Tensor
    x = torch.unsqueeze(x, 0).to(device)
    # model.train()
    # model(x)
    model.eval()
    # 出现  warnings.warn("Provided key {} for dynamic axes is not a valid input/output name".format(key)) 警告
    # 可能 是torch版本问题
    torch.onnx.export(model,
                      args=(x,),  # 输入数据 只用于执行一次模形，以便于获得模型的计算图
                      # args=(x1.cuda().half(), x2.cuda().half()),
                      f=onnx_filename,  # 导出的 文件名或 文件对象
                      export_params=True,  # 一般都导出训练后模型参数
                      verbose=False,  # 执行时是否显示过程信息
                      training=False,  # 是否用training模式导出模型
                      input_names=['input', ],  # 输入的口的名称 可存在多个输入口
                      output_names=['output', ],  # 输出的口的名称 可存在多个输出口
                      opset_version=None,  # 可选的onnx opset的版本 默认是torch 支持的最新版
                      do_constant_folding=False,  # 模型在导出过程中是否采用常数折叠优化。
                      # 常数折叠优化将用预先计算的常量节点替换一些具有所有常量输入的操作。
                      dynamic_axes=dynamic_axes  # 指定输入 和 输出 哪个维度值可变
                      # {'input_1':[0, 2, 3], 'input_2':[0], 'output':[0, 1]} list形式  前面的key对应上面xxx_names列表
                      # [0,1,...]对应第几维度值可变
                      # {'input_1':{0:'batch', 1:'width', 2:'height'}}dict形式 前面的key对应上面xxx_names列表
                      # {0:'batch', 1:'width', 2:'height'} 意思为 0,1,2维度可变  其名称对应为batch width height
                      )

    model = None

    #################### onnx检查生成模型 #######################
    checknet = onnx.load(onnx_filename)

    checknet: onnx.ModelProto
    onnx.checker.check_model(checknet)  # 检查文件模型是否正确
    print(onnx.helper.printable_graph(checknet.graph))  # 输出onnx的计算图
    # onnx.save_model(checknet, onnx_filename)
    checknet = None

    ########################测试生成的onnx文件###################
    x = dataset[0]
    x: np.ndarray
    cv.imshow("input", x)
    cv.waitKey(1)
    x = x.reshape(1, 1, *x.shape).astype('float32') / 255
    session = onnxruntime.InferenceSession(onnx_filename)  # 创建一个运行session，类似于tensorflow
    inputs_arg = {
        "input": x,
    }
    output = session.run(output_names=["output"],  # None,  # 过选的输出口名字 None表示全部输出
                         input_feed=inputs_arg  # 输入参数
                         )  # 模型运行，注意这里的输入必须是numpy类型

    y = output[0][0]  # output输出的是元组
    y = np.transpose(y, (1, 2, 0)) * 255
    y: np.ndarray
    y = y.astype(np.uint8)
    cv.imshow("output", y)
    key = cv.waitKey(0)
    if key == 27:
        exit()
    cv.destroyAllWindows()

    ########### netron网页显示 模型##################
    netron.start(onnx_filename)
